// 引入Node.js内置的文件系统模块，用于读取文件和目录操作
const fs = require('fs');
// 引入Node.js内置的路径模块，用于处理文件路径拼接和解析
const path = require('path');
// 引入Sharp第三方库，用于高性能图像处理、格式转换和像素数据提取
const sharp = require('sharp');

// 定义简单图像分类器类，实现基于特征提取的猫狗二分类识别
class SimpleImageClassifier {
  // 静态变量：控制训练时使用的图片数量，可通过修改此值调整训练规模
  static MAX_TRAINING_SAMPLES = 20000; // 默认200张图片，可修改为20000使用全部训练数据
  
  // 构造函数，初始化分类器实例的基本属性和参数
  constructor() {
    // 特征数组，存储所有训练样本提取的18维特征向量
    this.features = [];
    // 标签数组，存储对应训练样本的类别标签（0=猫，1=狗）
    this.labels = [];
    // 权重向量，存储逻辑回归模型的学习参数，训练前为null
    this.weights = null;
    // 图像统一尺寸，所有输入图像都会被调整为32x32像素以提高处理效率
    this.imageSize = 32; // 使用较小尺寸减少计算量，提高训练速度
  }

  // 从图像中提取增强特征的异步方法 - 修复版本
  async extractFeatures(imagePath) {
    try {
      // 使用Sharp库处理图像：调整尺寸并获取原始像素数据
      const { data, info } = await sharp(imagePath)
        .resize(this.imageSize, this.imageSize) // 调整图像为32x32像素
        .raw() // 获取原始RGB像素数据
        .toBuffer({ resolveWithObject: true }); // 返回包含数据和信息的对象
      
      // 初始化特征数组
      const features = [];
      // 计算总像素数量
      const pixelCount = this.imageSize * this.imageSize;
      
      // 初始化RGB三个通道的累加器和方差计算
      let rSum = 0, gSum = 0, bSum = 0;
      let rSquareSum = 0, gSquareSum = 0, bSquareSum = 0;
      
      // 遍历所有像素，每3个字节代表一个像素的RGB值
      for (let i = 0; i < data.length; i += 3) {
        const r = data[i];
        const g = data[i + 1];
        const b = data[i + 2];
        
        rSum += r;     // 累加红色通道值
        gSum += g;     // 累加绿色通道值
        bSum += b;     // 累加蓝色通道值
        
        // 累加平方值用于计算方差
        rSquareSum += r * r;
        gSquareSum += g * g;
        bSquareSum += b * b;
      }
      
      // 计算RGB通道的平均值并归一化到[0,1]区间
      const rMean = rSum / pixelCount / 255;
      const gMean = gSum / pixelCount / 255;
      const bMean = bSum / pixelCount / 255;
      
      features.push(rMean); // 红色通道平均值
      features.push(gMean); // 绿色通道平均值
      features.push(bMean); // 蓝色通道平均值
      
      // 计算RGB通道的方差（颜色分布特征）
      const rVar = (rSquareSum / pixelCount - Math.pow(rSum / pixelCount, 2)) / (255 * 255);
      const gVar = (gSquareSum / pixelCount - Math.pow(gSum / pixelCount, 2)) / (255 * 255);
      const bVar = (bSquareSum / pixelCount - Math.pow(bSum / pixelCount, 2)) / (255 * 255);
      
      features.push(rVar); // 红色通道方差
      features.push(gVar); // 绿色通道方差
      features.push(bVar); // 蓝色通道方差
      
      // 计算颜色比例特征（猫狗颜色差异）
      const totalColor = rSum + gSum + bSum;
      if (totalColor > 0) {
        features.push(rSum / totalColor); // 红色比例
        features.push(gSum / totalColor); // 绿色比例
        features.push(bSum / totalColor); // 蓝色比例
      } else {
        features.push(0.33, 0.33, 0.33); // 默认均匀分布
      }
      
      // 计算亮度分布特征（更细致的分级）
      let veryDarkPixels = 0, darkPixels = 0, midPixels = 0, brightPixels = 0, veryBrightPixels = 0;
      let edgePixels = 0; // 边缘像素计数
      
      // 遍历像素统计亮度分布和边缘特征
      for (let y = 0; y < this.imageSize; y++) {
        for (let x = 0; x < this.imageSize; x++) {
          const idx = (y * this.imageSize + x) * 3;
          const r = data[idx];
          const g = data[idx + 1];
          const b = data[idx + 2];
          
          // 计算当前像素的亮度值
          const brightness = (r + g + b) / 3;
          
          // 更细致的亮度分级
          if (brightness < 51) veryDarkPixels++;
          else if (brightness < 102) darkPixels++;
          else if (brightness < 153) midPixels++;
          else if (brightness < 204) brightPixels++;
          else veryBrightPixels++;
          
          // 简单边缘检测（与相邻像素的差异）
          if (x > 0 && y > 0) {
            const leftIdx = (y * this.imageSize + (x - 1)) * 3;
            const topIdx = ((y - 1) * this.imageSize + x) * 3;
            
            const leftBrightness = (data[leftIdx] + data[leftIdx + 1] + data[leftIdx + 2]) / 3;
            const topBrightness = (data[topIdx] + data[topIdx + 1] + data[topIdx + 2]) / 3;
            
            if (Math.abs(brightness - leftBrightness) > 30 || Math.abs(brightness - topBrightness) > 30) {
              edgePixels++;
            }
          }
        }
      }
      
      // 添加亮度分布特征
      features.push(veryDarkPixels / pixelCount);  // 很暗像素比例
      features.push(darkPixels / pixelCount);      // 暗像素比例
      features.push(midPixels / pixelCount);       // 中等亮度像素比例
      features.push(brightPixels / pixelCount);    // 亮像素比例
      features.push(veryBrightPixels / pixelCount); // 很亮像素比例
      features.push(edgePixels / pixelCount);      // 边缘像素比例
      
      // 计算图像对比度和纹理特征
      let contrast = 0;
      let texture = 0;
      const avgBrightness = (rSum + gSum + bSum) / (3 * pixelCount);
      
      // 计算对比度和纹理复杂度
      for (let i = 0; i < data.length; i += 3) {
        const brightness = (data[i] + data[i + 1] + data[i + 2]) / 3;
        const diff = Math.abs(brightness - avgBrightness);
        contrast += diff;
        texture += diff * diff; // 纹理复杂度
      }
      
      features.push(contrast / pixelCount / 255); // 对比度
      features.push(Math.sqrt(texture / pixelCount) / 255); // 纹理复杂度
      
      // 计算颜色饱和度特征
      let saturationSum = 0;
      for (let i = 0; i < data.length; i += 3) {
        const r = data[i] / 255;
        const g = data[i + 1] / 255;
        const b = data[i + 2] / 255;
        
        const max = Math.max(r, g, b);
        const min = Math.min(r, g, b);
        const saturation = max === 0 ? 0 : (max - min) / max;
        saturationSum += saturation;
      }
      features.push(saturationSum / pixelCount); // 平均饱和度
      
      // 返回增强的特征向量
      return features;
    } catch (error) {
      // 如果图像处理出错，打印错误信息并返回null
      console.error(`Error processing image ${imagePath}:`, error);
      return null;
    }
  }

  // 加载平衡训练数据的异步方法 - 修复版本
  async loadTrainingData(dataDir, maxSamples = SimpleImageClassifier.MAX_TRAINING_SAMPLES) {
    // 打印加载信息
    console.log(`Loading balanced training data from ${dataDir}...`);
    console.log(`🎯 目标：解决识别结果全是狗的问题`);
    
    // 检查目录是否存在
    if (!fs.existsSync(dataDir)) {
      throw new Error(`Training data directory does not exist: ${dataDir}`);
    }
    
    // 读取目录中的所有文件，分别统计猫狗图片
    const allFiles = fs.readdirSync(dataDir)
      .filter(file => file.toLowerCase().endsWith('.jpg'));
    
    if (allFiles.length === 0) {
      throw new Error(`No JPG files found in training directory: ${dataDir}`);
    }
    
    const catFiles = allFiles.filter(file => file.startsWith('cat.')).slice(0, Math.floor(maxSamples / 2));
    const dogFiles = allFiles.filter(file => file.startsWith('dog.')).slice(0, Math.floor(maxSamples / 2));
    
    console.log(`🐱 找到猫图片: ${catFiles.length} 张`);
    console.log(`🐶 找到狗图片: ${dogFiles.length} 张`);
    
    // 确保数据平衡
    const minCount = Math.min(catFiles.length, dogFiles.length);
    const balancedCatFiles = catFiles.slice(0, minCount);
    const balancedDogFiles = dogFiles.slice(0, minCount);
    
    console.log(`📊 使用平衡数据: 猫 ${balancedCatFiles.length} 张, 狗 ${balancedDogFiles.length} 张`);
    
    // 处理猫的图片
    console.log('🐱 正在处理猫的图片...');
    for (let i = 0; i < balancedCatFiles.length; i++) {
      const file = balancedCatFiles[i];
      const imagePath = path.join(dataDir, file);
      
      const features = await this.extractFeatures(imagePath);
      if (features) {
        this.features.push(features);
        this.labels.push(0); // 猫的标签为0
        
        if ((i + 1) % 50 === 0) {
          console.log(`  已处理猫图片 ${i + 1}/${balancedCatFiles.length} 张`);
        }
      }
    }
    
    // 处理狗的图片
    console.log('🐶 正在处理狗的图片...');
    for (let i = 0; i < balancedDogFiles.length; i++) {
      const file = balancedDogFiles[i];
      const imagePath = path.join(dataDir, file);
      
      const features = await this.extractFeatures(imagePath);
      if (features) {
        this.features.push(features);
        this.labels.push(1); // 狗的标签为1
        
        if ((i + 1) % 50 === 0) {
          console.log(`  已处理狗图片 ${i + 1}/${balancedDogFiles.length} 张`);
        }
      }
    }
    
    // 统计最终数据
    const catCount = this.labels.filter(label => label === 0).length;
    const dogCount = this.labels.filter(label => label === 1).length;
    
    console.log(`✅ 成功加载平衡数据: 猫 ${catCount} 张, 狗 ${dogCount} 张`);
    console.log(`📊 数据平衡比例: ${(catCount / (catCount + dogCount) * 100).toFixed(1)}% : ${(dogCount / (catCount + dogCount) * 100).toFixed(1)}%`);
    
    // 打乱数据顺序，避免训练偏向
    this.shuffleData();
  }
  
  // 打乱训练数据的方法
  shuffleData() {
    console.log('🔀 打乱训练数据顺序...');
    const indices = Array.from({length: this.features.length}, (_, i) => i);
    
    // Fisher-Yates 洗牌算法
    for (let i = indices.length - 1; i > 0; i--) {
      const j = Math.floor(Math.random() * (i + 1));
      [indices[i], indices[j]] = [indices[j], indices[i]];
    }
    
    // 重新排列特征和标签
    const shuffledFeatures = indices.map(i => this.features[i]);
    const shuffledLabels = indices.map(i => this.labels[i]);
    
    this.features = shuffledFeatures;
    this.labels = shuffledLabels;
  }

  // 使用改进的逻辑回归训练模型 - 修复版本
  train(epochs = 200, learningRate = 0.05) {
    console.log('🎓 开始训练改进版分类器...');
    console.log('🎯 目标：解决识别结果全是狗的问题');
    
    // 检查是否有训练数据
    if (!this.features || this.features.length === 0) {
      throw new Error('No training data loaded. Please call loadTrainingData() first.');
    }
    
    // 获取特征维度和样本数量
    const numFeatures = this.features[0].length;
    const numSamples = this.features.length;
    
    console.log(`📊 特征维度: ${numFeatures}`);
    console.log(`📊 训练样本: ${numSamples}`);
    
    // 改进的权重初始化 - 使用Xavier初始化
    this.weights = new Array(numFeatures + 1).fill(0).map(() => {
      return (Math.random() - 0.5) * 2 * Math.sqrt(6 / (numFeatures + 1));
    });
    
    // 统计类别分布
    const catCount = this.labels.filter(label => label === 0).length;
    const dogCount = this.labels.filter(label => label === 1).length;
    console.log(`📊 类别分布: 猫 ${catCount} 张, 狗 ${dogCount} 张`);
    
    // 计算类别权重，平衡训练
    const catWeight = numSamples / (2 * catCount);
    const dogWeight = numSamples / (2 * dogCount);
    console.log(`⚖️ 类别权重: 猫 ${catWeight.toFixed(2)}, 狗 ${dogWeight.toFixed(2)}`);
    
    let bestAccuracy = 0;
    let bestWeights = [...this.weights];
    
    // 开始训练循环
    for (let epoch = 0; epoch < epochs; epoch++) {
      let totalLoss = 0;
      let correct = 0;
      let catCorrect = 0, dogCorrect = 0;
      let catTotal = 0, dogTotal = 0;
      
      // 动态学习率调整
      const currentLR = learningRate * Math.pow(0.95, Math.floor(epoch / 50));
      
      // 遍历所有训练样本
      for (let i = 0; i < numSamples; i++) {
        const features = [1, ...this.features[i]];
        const target = this.labels[i];
        
        // 前向传播
        let prediction = 0;
        for (let j = 0; j < features.length; j++) {
          prediction += this.weights[j] * features[j];
        }
        
        // Sigmoid激活函数
        const output = 1 / (1 + Math.exp(-Math.max(-500, Math.min(500, prediction))));
        
        // 计算加权损失
        const sampleWeight = target === 0 ? catWeight : dogWeight;
        const loss = sampleWeight * (target * Math.log(output + 1e-15) + (1 - target) * Math.log(1 - output + 1e-15));
        totalLoss -= loss;
        
        // 统计准确率
        const predictedClass = output > 0.5 ? 1 : 0;
        if (predictedClass === target) {
          correct++;
          if (target === 0) catCorrect++;
          else dogCorrect++;
        }
        
        if (target === 0) catTotal++;
        else dogTotal++;
        
        // 反向传播 - 使用加权梯度
        const error = (output - target) * sampleWeight;
        for (let j = 0; j < features.length; j++) {
          this.weights[j] -= currentLR * error * features[j];
          
          // L2正则化防止过拟合
          if (j > 0) { // 不对偏置项进行正则化
            this.weights[j] -= currentLR * 0.001 * this.weights[j];
          }
        }
      }
      
      const accuracy = correct / numSamples;
      const catAccuracy = catTotal > 0 ? catCorrect / catTotal : 0;
      const dogAccuracy = dogTotal > 0 ? dogCorrect / dogTotal : 0;
      
      // 保存最佳权重
      if (accuracy > bestAccuracy) {
        bestAccuracy = accuracy;
        bestWeights = [...this.weights];
      }
      
      // 打印训练进度
      if ((epoch + 1) % 20 === 0 || epoch === 0) {
        const avgLoss = totalLoss / numSamples;
        console.log(`Epoch ${epoch + 1}/${epochs}`);
        console.log(`  总体准确率: ${(accuracy * 100).toFixed(2)}%`);
        console.log(`  猫识别率: ${(catAccuracy * 100).toFixed(2)}% (${catCorrect}/${catTotal})`);
        console.log(`  狗识别率: ${(dogAccuracy * 100).toFixed(2)}% (${dogCorrect}/${dogTotal})`);
        console.log(`  损失: ${avgLoss.toFixed(4)}, 学习率: ${currentLR.toFixed(4)}`);
        
        // 检查是否解决了"全是狗"问题
        if (catCorrect > 0) {
          console.log(`  ✅ 已能识别出猫！`);
        } else {
          console.log(`  ⚠️ 仍然无法识别猫`);
        }
        console.log('');
      }
    }
    
    // 使用最佳权重
    this.weights = bestWeights;
    
    console.log('🎉 训练完成！');
    console.log(`🏆 最佳准确率: ${(bestAccuracy * 100).toFixed(2)}%`);
    
    // 最终验证
    this.validateTraining();
  }
  
  // 验证训练效果的方法
  validateTraining() {
    console.log('🔍 验证训练效果...');
    
    let catCorrect = 0, dogCorrect = 0;
    let catTotal = 0, dogTotal = 0;
    let catPredictions = 0, dogPredictions = 0;
    
    for (let i = 0; i < this.features.length; i++) {
      const features = [1, ...this.features[i]];
      const target = this.labels[i];
      
      let prediction = 0;
      for (let j = 0; j < features.length; j++) {
        prediction += this.weights[j] * features[j];
      }
      
      const output = 1 / (1 + Math.exp(-prediction));
      const predictedClass = output > 0.5 ? 1 : 0;
      
      if (predictedClass === 0) catPredictions++;
      else dogPredictions++;
      
      if (target === 0) {
        catTotal++;
        if (predictedClass === 0) catCorrect++;
      } else {
        dogTotal++;
        if (predictedClass === 1) dogCorrect++;
      }
    }
    
    console.log('📊 训练集验证结果:');
    console.log(`  预测分布: 猫 ${catPredictions} 张, 狗 ${dogPredictions} 张`);
    console.log(`  猫识别准确率: ${(catCorrect / catTotal * 100).toFixed(1)}%`);
    console.log(`  狗识别准确率: ${(dogCorrect / dogTotal * 100).toFixed(1)}%`);
    
    if (catPredictions > 0) {
      console.log('✅ 成功解决"全是狗"的问题！');
    } else {
      console.log('❌ 仍然存在"全是狗"的问题');
    }
  }

  // 预测单张图像类别的异步方法 - 修复版本
  async predict(imagePath) {
    // 检查模型是否已经训练
    if (!this.weights) {
      throw new Error('Model not trained yet');
    }
    
    // 提取待预测图像的特征
    const features = await this.extractFeatures(imagePath);
    if (!features) {
      return null;
    }
    
    // 构建输入向量（添加偏置项）
    const input = [1, ...features];
    let prediction = 0;
    
    // 计算线性组合 w·x + b
    for (let i = 0; i < input.length; i++) {
      prediction += this.weights[i] * input[i];
    }
    
    // 通过Sigmoid函数得到概率，限制范围防止溢出
    const clampedPrediction = Math.max(-500, Math.min(500, prediction));
    const dogProbability = 1 / (1 + Math.exp(-clampedPrediction));
    const catProbability = 1 - dogProbability;
    
    // 使用更平衡的决策阈值
    const threshold = 0.5;
    const predictedClass = dogProbability > threshold ? 1 : 0;
    const confidence = Math.max(dogProbability, catProbability);
    
    // 返回详细的预测结果
    return {
      class: predictedClass,
      className: predictedClass === 0 ? 'cat' : 'dog',
      confidence: confidence,
      probability: predictedClass === 1 ? dogProbability : catProbability,
      probabilities: {
        cat: catProbability,
        dog: dogProbability
      },
      rawPrediction: prediction,
      features: {
        avgBrightness: features[0] + features[1] + features[2], // RGB平均值
        colorVariance: features[3] + features[4] + features[5], // RGB方差
        edgeDensity: features[features.length - 3], // 边缘密度
        textureComplexity: features[features.length - 2], // 纹理复杂度
        saturation: features[features.length - 1] // 饱和度
      }
    };
  }

  // 保存训练好的模型到文件
  saveModel(filePath) {
    // 构建模型数据对象
    const modelData = {
      weights: this.weights,     // 权重向量
      imageSize: this.imageSize  // 图像尺寸参数
    };
    // 将模型数据写入JSON文件
    fs.writeFileSync(filePath, JSON.stringify(modelData, null, 2));
    console.log(`Model saved to ${filePath}`); // 打印保存成功信息
  }

  // 从文件加载训练好的模型
  loadModel(filePath) {
    // 读取并解析JSON模型文件
    const modelData = JSON.parse(fs.readFileSync(filePath, 'utf8'));
    // 恢复权重向量
    this.weights = modelData.weights;
    // 恢复图像尺寸参数
    this.imageSize = modelData.imageSize;
    console.log(`Model loaded from ${filePath}`); // 打印加载成功信息
  }
}

// 主函数，演示分类器的完整使用流程
async function main() {
  console.log('=== 简单猫狗分类器 ===');
  
  try {
    // 创建分类器实例
    const classifier = new SimpleImageClassifier();
    
    // 加载训练数据并训练模型
    const trainDataPath = path.resolve(__dirname, '../../images/train');
    await classifier.loadTrainingData(trainDataPath, 200);
    classifier.train(200, 0.1); // 训练200个epoch，学习率0.1
    
    // 保存训练好的模型
    classifier.saveModel(path.join(__dirname, '../../models/simple_model.json'));
    
    // 测试单张图片预测
    console.log('\n测试预测:');
    const testImagePath = path.join(__dirname, '../../images/test1/1.jpg');
    
    // 检查测试图片是否存在
    if (fs.existsSync(testImagePath)) {
      const result = await classifier.predict(testImagePath); // 进行预测
      if (result) {
        // 打印预测结果
        console.log(`预测结果: ${result.className}`);
        console.log(`置信度: ${(result.probability * 100).toFixed(2)}%`);
        console.log(`详细概率: 猫 ${(result.probabilities.cat * 100).toFixed(2)}%, 狗 ${(result.probabilities.dog * 100).toFixed(2)}%`);
      }
    }
    
    // 批量测试前10张图片
    console.log('\n批量测试前10张图片:');
    const testDir = path.join(__dirname, '../../images/test1'); // 测试目录路径
    // 获取测试目录中的前10个jpg文件
    const testFiles = fs.readdirSync(testDir)
      .filter(file => file.endsWith('.jpg')) // 过滤jpg文件
      .slice(0, 10); // 只取前10个
    
    // 遍历测试文件进行预测
    for (const file of testFiles) {
      const testPath = path.join(testDir, file); // 构建完整路径
      const result = await classifier.predict(testPath); // 预测
      if (result) {
        // 打印每个文件的预测结果
        console.log(`${file}: ${result.className} (${(result.probability * 100).toFixed(1)}%)`);
      }
    }
    
  } catch (error) {
    // 捕获并打印任何错误
    console.error('Error:', error);
  }
}

// 如果直接运行此文件，则执行主函数
if (require.main === module) {
  main();
}

// 导出SimpleImageClassifier类供其他模块使用
module.exports = SimpleImageClassifier;